#pragma once

#include "txdnn_internal.h"
#include "operators/rope/operator_description.h"
#include "operators/operator_base.h"

#include "backend/kernel_interface/kernel_build_params.h"
// #include "backend/kernel_interface/kernel_scheduler.h"
#include "operators/invoker_params_base.h"


namespace txdnn {
namespace operators {
namespace rope {

using RoPEOperatorFwd = NonTunableOperatorBase<txdnn::ExecutionContext, txdnn::operators::rope::OperatorDescriptionFwd>;

struct RoPEForward final : RoPEOperatorFwd {

  Operators getOperator(const txdnn::ExecutionContext& ctx, 
                        const txdnn::operators::rope::OperatorDescriptionFwd& op_desc) const override {
    auto ops = Operators{};

    auto dtype          = op_desc.getXDesc().getDataType();
    auto input_dtype    = GetDataType(dtype);
    auto output_dtype   = GetDataType(op_desc.getYDesc().getDataType());
    auto y_dims         = op_desc.getYDesc().getDims();
    auto output_num_el   = std::accumulate(y_dims.begin(), y_dims.end(), 1ULL, std::multiplies<size_t>{});

    {
      auto kernel = KernelInfo{};
      kernel.kernel_path = "TxDNN_RoPE.c";  // kernel file
      kernel.kernel_name = "RoPEFwd";       // kernel function name

      const auto build_params = txdnn::backend::KernelBuildParams{}; // TODO set build options when initialization
      kernel.compile_options = build_params.generatorFor(txdnn::backend::kbp::MainC{});

      ops.construct_params_.emplace_back(kernel);
    }

    ops.invoker_factory_ = [](const std::vector<txdnn::backend::Kernel>& kernels) {
      return [=](const txdnn::Handle& handle, const txdnn::AnyInvokeParams& raw_params) {
        auto kernel = handle.run(kernels.front());

        // real kernel call 
        // kernel();
      };
    };
    return ops;
  }


};

} // namespace rope
} // namespace operators
} // namespace txdnn